Introduction à l'apprentissage par renforcement¶
TP 1 - les manchots multi-bras¶
1/4 de la note finale est liée à la mise en forme :
- pensez à nettoyer les outputs inutiles (installation, messages de débuggage, ...)
- soignez vos figures : les axes sont-ils faciles à comprendre ? L'échelle est adaptée ?
- commentez vos résultats : vous attendiez-vous à les avoir ? Est-ce étonnant ? Faites le lien avec la théorie.
Ce TP reprend l'exemple d'un médecin et de ses vaccins. Vous allez comparer plusieurs stratégies et trouver celle optimale. Un TP se fait seul ou en binôme. Aucun groupe de plus de 2 personnes.
Vous allez rendre le TP depuis un lien GitHub avec ce notebook mais une version du rapport exportée en PDF & HTML.
! pip install matplotlib tqdm numpy ipympl opencv-python torch seaborn
!jupyter labextension install @jupyter-widgets/jupyterlab-manager
!jupyter labextension install jupyter-matplotlib
%load_ext autoreload
%autoreload 2
%matplotlib inline
import typing as t
import math
import torch
import numpy as np
from tqdm.auto import trange, tqdm
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.animation as animation
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import cv2
from IPython.display import display, clear_output, HTML
import seaborn as sns
torch.random.manual_seed(0)
K = 5 # num arms
Présentation du problème¶
class ArmBernoulli:
def __init__(self, p: float):
"""
Vaccine treatment following a Bernoulli law (mean is p and variance is p(1-p)
Args:
p (float): mean parameter
>>> torch.random.manual_seed(random_state)
>>> arm = ArmBernoulli(0.5)
>>> arm.sample(5)
tensor([ True, False, True, True, True])
"""
self.immunity_rate = p
def sample(self, n: int = 1):
return torch.rand(n) < self.immunity_rate
def __repr__(self):
return f'<ArmBernoulli p={self.immunity_rate}'
def generate_arms(num_arms: int):
means = torch.rand(num_arms)
MAB = [ArmBernoulli(m) for m in means]
assert MAB[0].immunity_rate == means[0]
assert (MAB[0].sample(10) <= 1).all() and (MAB[0].sample(10) >= 0).all()
return MAB
MAB = generate_arms(K)
Ce TP reprend l'exemple du médecin présenté en cours.
Q1. Créez une fonction pour trouver $\mu^*$ à partir d'un MAB. Comment est définie la récompense $R_k$ ? Que représente concrètement le regret dans le contexte de ce TP ?
def best_vaccine_in_mab(mab):
# return the index of the best vaccine in the MAB
return np.array([vaccine.immunity_rate for vaccine in mab]).argmax()
print(f"Best vaccine in MAB: {best_vaccine_in_mab(MAB)}, with immunity rate {MAB[best_vaccine_in_mab(MAB)].immunity_rate.item():.2f}")
for i, arm in enumerate(MAB):
print(f"Vaccine {i} has immunity rate {arm.immunity_rate.item():.2f}")
Best vaccine in MAB: 1, with immunity rate 0.77 Vaccine 0 has immunity rate 0.50 Vaccine 1 has immunity rate 0.77 Vaccine 2 has immunity rate 0.09 Vaccine 3 has immunity rate 0.13 Vaccine 4 has immunity rate 0.31
La récompense $R_k$ est définie selon la loi de Bernoulli, donc 1 (p) si c'est une réussite et 0 (1-p) si échec.
Le regret représente la différence entre le fonctionnement théorique du vaccin (phase de d'exploration) et l'exploitation réelle de la chance d'immunisation.
Note importante : pour la suite, les résultats seront généralement réalisés avec 100 initialisations différentes du MAB (tous les MAB ont 5 vaccins mais des taux d'immunistation différent) pour réduire le bruit de simulation. Concrètement, on exécutera au moins 100x generate_arms.
I. Cas classique des bandits manchots¶
I.a. Solution Gloutonne¶
Le médecin fonctionne sur deux phases :
- Exploration : Le médecin calcule le taux d'immunisation empirique sur les N premiers patients en administrant le même nombre de fois chaque vaccin :
avec $T_i = \sum_{k=0}^{N-1} \chi_{v_k,i}$.
- Exploitation : Le vaccin $v_i = \arg\max_j \hat{\mu_j}[0\rightarrow N]$ est utilisé pour les M patients suivants.
Q2. Implémentez cette solution avec N = 50 et M = 500 et testez-la avec 100 MAB. On souhaite savoir si vous trouvez le vaccin optimal à l'issue d'une phase d'exploration. Quelle est l'espérance empirique de cette variable ? Et son écart-type ? Calculez de même l'espérance et l'écart-type du regret sur vos 100 simulations.
Pour rappel, le regret est défini par :
$$r_n = n\mu^* - \sum_{k=0}^{n-1} R_k$$Attention : $n$ est le nombre total de patients, donc ici $N + M$.
Mab_list = [generate_arms(5) for _ in range(100)]
N = 50
M = 500
def exploration(n, mab):
R = np.zeros(K)
for i, arm in enumerate(mab):
R[i] = arm.sample(n // K).sum()
return R
def getBestVaccin(R, hist):
Ri = np.array([np.inf] * K) # Important pour n = 0
for i in range(K):
if hist[i] != 0:
Ri[i] = R[i] / hist[i]
return Ri.argmax()
def exploitationQ2(n,m, mab_list):
rates = []
regrets = []
for mab in mab_list:
# Phase d'exploration
R = exploration(n, mab)
# exploitation
hist = [n // K] * K
bestVac= getBestVaccin(R, hist)
Mab_res = mab[bestVac].sample(m).sum()
R[bestVac] += Mab_res
hist[bestVac] += m
### exploration success rate
list_immunity_rates = [i.immunity_rate for i in mab]
best_vac_from_immunity_rates = best_vaccine_in_mab(mab)
success = (bestVac == best_vac_from_immunity_rates)
### regret
regret = (n + m) * max(list_immunity_rates) - R.sum()
regrets.append(regret)
rates.append(success)
return regrets, rates
regrets, rates = exploitationQ2(N, M, Mab_list)
mean_rates = np.mean(rates)
std_rates = np.std(rates)
mean_regret = np.mean(regrets)
std_regret = np.std(regrets)
print(f"Espérance empirique: {mean_rates}")
print(f"Écart-type : {std_rates}")
print(f"Espérance du regret: {mean_regret}")
print(f"Écart-type du regret: {std_regret}")
Espérance empirique: 0.68 Écart-type : 0.46647615158762396 Espérance du regret: 30.894350051879883 Écart-type du regret: 31.458023071289062
On remarque que l'algorithme glouton ne trouve pas tout le temps le bon vaccin, cependant l'espérance du regret reste aux alentours de 30. On peut en conclure que l'algorithme glouton ne trouve pas forcément le meilleur vaccin, mais il en trouvera un qui marche bien ou qui est proche du taux d'immunité du meilleur vaccin. Cette théorie peut être soutenue par un écart-type relativement petit; il serait bien plus grand si l'algorithme ne trouvait pas un vaccin presque aussi efficace que le meilleur vaccin.
Q3. On étudie maintenant l'influence de la taille du training set $N$. On considère que N+M est une constante, puis on fait varier N entre K et M. Calculez le regret pour ces différentes tailles du training set différents MAB et representez le regret moyen, le regret min et max (vous devriez trouver une courbe en U ou en V pour le regret moyen). Quelle est la taille optimale du training set ?
def exploitationQ3(n,m, mab_list):
regrets = []
for mab in mab_list:
# Phase d'exploration
R = exploration(n, mab)
# exploitation
hist = [n // K] * K
bestVac= getBestVaccin(R, hist)
Mab_res = mab[bestVac].sample(m).sum()
R[bestVac] += Mab_res
hist[bestVac] += m
### regret
regret = (n + m) * max([i.immunity_rate for i in mab]) - R.sum()
regrets.append(regret)
return regrets
m = N + M
Mab_list_regrets = []
for n in range(K,m, K):
regrets = exploitationQ3(n, m - n, Mab_list)
Mab_list_regrets.append(np.mean(regrets))
best_N = int(min(Mab_list_regrets))
best_M = int(m - best_N)
print(f'min: {min(Mab_list_regrets):.2f}, k: {Mab_list_regrets.index(min(Mab_list_regrets)) * 5}')
print(f'max: {max(Mab_list_regrets):.2f}, k: {Mab_list_regrets.index(max(Mab_list_regrets)) * 5}')
plt.plot(range(K,m, K), Mab_list_regrets, label='regret')
plt.xticks(np.arange(0, m, 25), rotation=-90)
plt.title("Regret moyen pour différentes tailles du training set")
plt.xlabel("Taille du training set N")
plt.ylabel("Regret moyen")
plt.legend()
plt.show()
min: 25.13, k: 40 max: 189.38, k: 540
En lançant plusieurs fois l'algorithme, on remarque que la taille optimale de l'ensemble d'entraînement se trouve entre 30 et 60. Le graphique le confirme.
Avoir une phase d'exploration trop grande va réduire l'erreur possible due à l'aléa des MAB. Cependant, utiliser trop souvent de mauvais vaccins va réduire le nombre de patients soignés et, par la même occasion, les récompenses cumulées. Cela explique l'augmentation du regret à partir de 60.
Avant 30, l'algorithme n'a pas le temps de faire assez de tests sur chaque vaccin pour éviter une erreur d'approximation du taux d'immunité. Cette erreur vient de l'aléa des MAB. Ainsi, en sortant de la phase d'exploration, il se peut que le vaccin avec les meilleurs résultats soit le pire.
Q4. On propose d'améliorer l'algorithme précédant en mettant à jour les taux d'immunisation empiriques $\hat{\mu}_i$ pendant la phase d'exploitation (algorithme greedy). Concrètement, à chaque nouvel patient, on lui administre le meilleur vaccin selon les stats. Notez vous une amélioration du regret ? Proposez un exemple où les taux d'immunisation du MAB ne changent rien.
def exploitationQ4(n,m, mab_list):
regrets = []
for mab in mab_list:
# Phase d'exploration
R = exploration(n, mab)
# exploitation
list_immunity_rate = [i.immunity_rate.item() for i in mab]
bestVacGt = list_immunity_rate.index(max(list_immunity_rate))
hist = [n // K] * K
total_rewards = R.sum()
for _ in range(1, m + 1):
bestVac = getBestVaccin(R, hist)
Mab_res = mab[bestVac].sample()
hist[bestVac] += 1
R[bestVac] += Mab_res
total_rewards += Mab_res
# Mise à jour du regret à chaque étape
regret = (n + m) * list_immunity_rate[bestVacGt] - total_rewards
regrets.append(regret)
return regrets
regrets = exploitationQ4(N, M, Mab_list)
print(f'N = {N}, M = {M} regret mean: {np.mean(regrets):.2f}, regret standard deviation: {np.std(regrets):.2f}')
N = 50, M = 500 regret mean: 20.54, regret standard deviation: 11.59
On remarque une amélioration du regret par rapport à la question 3. L'écart-type a lui aussi largement baissé. Cette baisse, due à l'algorithme greedy, montre que même si la phase d'exploration ne trouve pas le meilleur vaccin, l'algorithme va se rattraper pendant la phase d'exploitation en changeant la stratégie de vaccination en fonction des nouvelles statistiques.
Dans le cas où, dès la phase d'exploration, on trouve directement le meilleur vaccin théorique, cette amélioration n'aura pas d'impact.
Q5. Nouvelle amélioration : à chaque nouveau patient, on choisit si on lui administre le meilleur vaccin avec une probabilité $\epsilon$ ou un vaccin aléatoire ($p=1-\epsilon$). Vérifiez si vous obtenez un meilleur résultat avec N = 0 ou N > 0. À votre avis, à quoi sert $\epsilon$ ?
def exploitationQ5(n,m, mab_list, epsilon):
regrets = []
list_regrets = []
for mab in mab_list:
# Phase d'exploration
R = exploration(n, mab)
# exploitation
list_immunity_rate = [i.immunity_rate.item() for i in mab]
bestVacGt = list_immunity_rate.index(max(list_immunity_rate))
hist = [n // K] * K
list = [(i * list_immunity_rate[bestVacGt] - R[:i].sum() ) for i in range(1, n + 1)]
for i in range(1, m + 1):
if (np.random.rand() <= epsilon):
bestVac = getBestVaccin(R, hist)
else:
bestVac = np.random.randint(K)
Mab_res = mab[bestVac].sample()
hist[bestVac] += 1
R[bestVac] += Mab_res
list.append((n + i) * list_immunity_rate[bestVacGt] - R.sum())
regrets.append(list[-1])
list_regrets.append(list)
return regrets, list_regrets
Mab_list_regrets = []
regrets_0, _ = exploitationQ5(0, M + N, Mab_list, 0.7)
regrets_N, list_regrets = exploitationQ5(N, M, Mab_list, 0.7)
print(f'N = 0, regret = {np.mean(regrets_0):.2f}')
print(f'N > 0, regret = {np.mean(regrets_N):.2f}')
N = 0, regret = 60.40 N > 0, regret = 68.10
On a choisit arbitrairement la valeur 0.7 pour epsilon.
epsilon sert à simuler une phase d'exploration avec N=1. Cela permet de tester des vaccins moins efficaces selon les stats, par moment un manque de chance sur le vaccin joue en ca défaveur, même si ce dernier à un taux d'immunité élevé.
On remarque une augmentation du regret par rapport aux questions précédentes. Cette augmentation est provoquée par l'ajout de l'amélioration. En effet, contrairement aux questions précédentes, l'algorithme utilisera potentiellement un vaccin très mauvais durant la phase d'exploitation. Là où précédemment on était plus ou moins assuré que le vaccin choisi était le meilleur ou proche du vaccin avec le plus haut taux d'immunité.
I.b. Borne inférieure de Lai & Robbins [Lai et Robbins, 1985]¶
Lai et Robbins [Lai et Robbins, 1985] considère une classe d'algorithmes $\pi$ pour résoudre ce type de problèmes.
Ils ont trouvé une borne inférieure sur les récompenses cumulées en valeur asymptotique :
$$\lim_{n\rightarrow \infty} \inf_{\pi} \frac{\sum_{k=0}^{n-1} R_k}{\log n} \geq \sum_{i~\text{tel que}~\mu_i \lt \mu^*} \frac{\mu^∗−\mu_i}{\text{KL}(\mu_i, \mu^*)} :=C(\mu)$$avec $\text{KL}(x, y) = x \log(x/y) + (1 − x) \log((1 − x)/(1 − y))$ (distance de Kullback-Leibler) et $\sum_{k=0}^{n-1} R_k$ la récompense obtenue sur $n$ patients.
Q6. Justifiez pourquoi on peut en déduire que le regret d'un algorithme raisonnable sera au pire logarithmique.
on remarque que $C(\mu)$ est une constante. cela veut dire que les récompenses cumulées croît au moins proportionnelement à $\log n$
En utilisant la borne inférieur sur la formule du regret $R(n) = n \mu^* - \sum_{k=0}^{n-1} R_k$ vu en cours on obtient : $$R(n) \leq n \mu^* - C(\mu)\log n$$
le regret est donc majoré par une expression linéaire $n - \log n$
Q7. Tracez le regret issu de la borne de Lai & Robbins et comparez le au regret obtenu avec l'algorithme glouton.
# nouvelle exploration qui stocke les regrets
def explorationQ7(n, mab):
R = np.zeros(K)
regrets = []
bv = np.max([i.immunity_rate for i in mab])
for i in range(n // K):
for j, arm in enumerate(mab):
R[j] += arm.sample()
regrets.append((K * i + (j + 1)) * bv - R.sum())
return R, regrets
# exploitation de la Q4 (sans epsilon)
def exploitationQ7(n,m, mab_list):
MABS_regrets = []
for mab in mab_list:
MAB_regrets = []
# Phase d'exploration
R, exploration_regrets = explorationQ7(n, mab)
MAB_regrets.extend(exploration_regrets)
# exploitation
bestVacGt = np.max([i.immunity_rate for i in mab])
hist = [n // K] * K
for i in range(1, m + 1):
bestVac = getBestVaccin(R, hist)
R[bestVac] += mab[bestVac].sample()
hist[bestVac] += 1
# Mise à jour du regret à chaque étape
regret = (n + i) * bestVacGt - R.sum()
MAB_regrets.append(regret)
MABS_regrets.append(MAB_regrets)
return MABS_regrets
def kl(x,y):
return x * np.log(x/y) + (1-x) * np.log((1-x)/(1-y))
def lower_bound(n, mab):
bestVacGt = max([mu.immunity_rate.item() for mu in mab])
C = 0
for mu in mab:
if mu.immunity_rate.item() == bestVacGt:
continue
C += (bestVacGt - mu.immunity_rate.item()) / kl(mu.immunity_rate.item(), bestVacGt)
return C * np.log(n)
list_c = []
for mab in Mab_list:
c = []
for i in range(1, N + M+1):
c.append(lower_bound(i, mab))
list_c.append(c)
regrets = exploitationQ7(N, M, Mab_list)
plt.plot(range(1, N + M+1), np.mean(list_c, axis=0), label='regret lower bound')
plt.plot(range(1, N + M+1), np.mean(regrets, axis=0), label='Regret')
plt.xticks(np.arange(0, N+M, 25), rotation=-90)
plt.title("différence entre le regret du *greedy algorithm* et de la borne inférieure")
plt.xlabel("Nombre de patients traités")
plt.ylabel("Regret moyen")
plt.legend()
plt.show()
La courbe de regret reste bien en dessous de la courbe de regret de la borne inférieure. Cela confirme la réponse de la question 6. Les 50 premiers regrets correspondent à la phase d'exploration, cela explique sa forme différente.
I.c. Upper Confidence Bounds¶
Cet algorithme améliore la version précédente en ajoutant un biais lié à la fréquentation de chaque vaccin :
$\bar{\mu}_i = \hat{\mu}_i + \sqrt{\frac{C\log{n}}{T_i}}$,
avec $C=2$.
Q8. Implémentez la modification de cet algorithme. Observez un intérêt à conserver $N > 0$ ? Et $\epsilon < 1$ ? Expliquez pourquoi.
Dans la suite, on prendra $N = 0$ et $\epsilon = 1$.
def getBestVaccinWithBiais(R, hist,nb, C=2):
Ri = np.zeros(K)
for i in range(K):
if hist[i] != 0:
Ri[i] = R[i] / hist[i] + math.sqrt(C * math.log(nb) / hist[i])
else:
Ri[i] = np.inf
return Ri.argmax()
def exploitationQ8(n,m, mab_list, epsilon):
regrets = []
for mab in mab_list:
# Phase d'exploration
R = exploration(n, mab)
# exploitation
list_immunity_rate = [i.immunity_rate.item() for i in mab]
bestVacGt = list_immunity_rate.index(max(list_immunity_rate))
hist = [n // K] * K
total_rewards = R.sum()
for i in range(1, m + 1):
if (np.random.rand() <= epsilon):
bestVac = getBestVaccinWithBiais(R, hist, n+i)
else:
bestVac = np.random.randint(K)
Mab_res = mab[bestVac].sample()
hist[bestVac] += 1
R[bestVac] += Mab_res
total_rewards += Mab_res
# Mise à jour du regret à chaque étape
regret = (n + i) * list_immunity_rate[bestVacGt] - total_rewards
regrets.append(regret)
return regrets
Mab_list_regrets = []
# Initialisation des listes pour chaque epsilon
m = N + M
borneN = [0,80]
epsilons = [0.1, 0.3, 0.5, 0.75, 1]
n_values = range(borneN[0], borneN[1], K) # borneN et K doivent être définis
Mab_list_regrets = {epsilon: [] for epsilon in epsilons} # Dictionnaire pour chaque epsilon
for n in n_values:
for epsilon in epsilons:
regrets = exploitationQ8(n, m - n, Mab_list, epsilon)
Mab_list_regrets[epsilon].append(np.mean(regrets))
# Tracer les résultats
for epsilon in epsilons:
plt.plot(n_values, Mab_list_regrets[epsilon], marker='o', linestyle='-', label=f'epsilon={epsilon}')
plt.title('Évolution des regrets moyens pour différentes valeurs de epsilon')
plt.xlabel('Valeurs de N')
plt.ylabel('Regrets moyens')
plt.legend() # Ajoute une légende pour chaque courbe
plt.grid(True)
plt.show()
$N > 0$ indique qu'il y aura une phase d'exploration avant la phase d'exploitation. Ainsi, avant d'exploiter les vaccins, on aura des statistiques plus ou moins bonnes sur le taux d'immunité. Or, la formule avec le biais de UCB force l'algorithme à tester des vaccins dont on n'a pas d'information. Sur un nombre de patients assez élevé, le biais rend la phase d'exploration inutile car chaque vaccin sera forcément testé.
$\epsilon < 1$ indique que pendant la phase d'exploitation, il y a une chance qu'on choisisse aléatoirement un vaccin, ce qui permet de rectifier une possible erreur due à un manque de chance dans la phase d'exploration. Or, comme expliqué plus haut, le biais s'assure que l'algorithme n'utilise pas tout le temps le même vaccin (pour un grand nombre de patients testés).
Le biais ajouté remplace donc ces deux paramètres, les rendant inutiles.
Q9. Tracez sous la forme d'une animation l'évolution du regret et l'évolution des taux d'immunisation empirique. Dans la figure de gauche, vous representerez $\bar{\mu}_i$ et $\hat{\mu}_i$.
def compute_erc(R, hist):
Ri = np.array([np.inf] * K) # usefull when n = 0
for i in range(K):
if hist[i] != 0:
Ri[i] = R[i] / hist[i]
return Ri
def compute_er_(R, hist, nb, C=2):
Ri = np.array([np.inf] * K) # usefull when n = 0
for i in range(K):
if hist[i] != 0:
Ri[i] = R[i] / hist[i] + math.sqrt(C * math.log(nb) / hist[i])
return Ri
def exploitationQ9(m, mab_list, C=2):
regrets = []
mu_ = []
muc = []
for mab in mab_list:
regret_list = []
rate_list = []
rate_list_ = []
# Phase d'exploration
R = np.zeros(K)
# Phase exploitation
bestVacGt = max([i.immunity_rate.item() for i in mab])
hist = np.zeros(K)
for i in range(1, m + 1):
bestVac = np.argmax(compute_er_(R, hist, i - 1, C))
# Mise à jour des résultats et de l'histogramme
R[bestVac] += mab[bestVac].sample()
hist[bestVac] += 1
# Mise à jour du regret à chaque étape
regret = (i) * bestVacGt - R.sum()
regret_list.append(regret)
# calcul de mû
erc = compute_erc(R, hist)
rate_list.append(erc)
# calcul de mu_
er_ = compute_er_(R, hist,i, C)
rate_list_.append(er_)
regrets.append(regret_list)
muc.append(rate_list)
mu_.append(rate_list_)
return regrets, muc, mu_
Mab_list_1 = [generate_arms(K) for _ in range(1)]
Mab_list_regrets = []
M = 550
# Run the exploitation function
regrets, muc, mu_ = exploitationQ9(M, Mab_list_1)
# Mean cumulative regret over simulations
mean_regret = np.mean(regrets, axis=0) # Shape: (M,)
# Mean muc and mu_ over simulations
mean_muc = np.mean(muc, axis=0) # Shape: (M, K)
mean_mu_ = np.mean(mu_, axis=0) # Shape: (M, K)
# Set up the figure and axes
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
# Initialize lines for the left plot
lines_muc = [ax1.plot([], [], label=f'$\overline{{\mu}}_{i+1}$')[0] for i in range(K)]
lines_mu_ = [ax1.plot([], [], linestyle='--', label=f'$\hat{{\mu}}_{i+1}$')[0] for i in range(K)]
# Initialize line for the right plot
line_regret, = ax2.plot([], [], label='Cumulative Regret')
# Set up the axes
ax1.set_xlim(0, M)
ax1.set_ylim(0, 2)
ax1.set_xlabel('Time Steps')
ax1.set_ylabel('Immunization Rate')
ax1.set_title('Evolution of Empirical Immunization Rates')
ax1.legend(loc='upper right')
ax1.set_xticks(np.arange(0, M, 50))
ax2.set_xlim(0, M)
ax2.set_ylim(0, np.max(mean_regret))
ax2.set_xlabel('Time Steps')
ax2.set_ylabel('Cumulative Regret')
ax2.set_title('Evolution of Cumulative Regret')
frames = np.arange(0, M, 5)
# Animation function
def animate(i):
x = np.arange(i)
# Update left plot with $\bar{\mu}_i$ and $\hat{\mu}_i$
for k in range(K):
# Update empirical mean immunization rate lines
y_muc = mean_muc[:i, k]
lines_muc[k].set_data(x, y_muc)
# Update upper confidence bound lines
y_mu_ = mean_mu_[:i, k]
lines_mu_[k].set_data(x, y_mu_)
# Update right plot with cumulative regret
y_regret = mean_regret[:i]
line_regret.set_data(x, y_regret)
return lines_muc + lines_mu_ + [line_regret]
# Create the animation
ani = animation.FuncAnimation(fig, animate, frames=frames, interval=50, blit=True, cache_frame_data=False)
HTML(ani.to_jshtml())